#!/usr/bin/env python3
"""
数据迁移脚本：将数据库中的Base64图片数据保存为文件
"""

import os
import base64
import sqlite3
from datetime import datetime

def migrate_images():
    """迁移图片数据"""
    db_path = 'instance/prompt_manager.db'
    
    if not os.path.exists(db_path):
        print(f"数据库文件不存在: {db_path}")
        return
    
    # 确保files目录存在
    os.makedirs('files', exist_ok=True)
    
    # 连接数据库
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    
    try:
        # 获取所有有图片数据的记录
        cursor.execute("SELECT id, chrome_id, image_data FROM prompts WHERE image_data IS NOT NULL AND image_data != ''")
        records = cursor.fetchall()
        
        print(f"找到 {len(records)} 条包含图片数据的记录")
        
        migrated_count = 0
        error_count = 0
        
        for record_id, chrome_id, image_data in records:
            try:
                if not image_data or not image_data.startswith('data:image'):
                    print(f"记录 {record_id} 的图片数据格式不正确，跳过")
                    continue
                
                # 解析Base64数据
                if ',' in image_data:
                    header, data = image_data.split(',', 1)
                    # 获取图片格式
                    if 'png' in header:
                        ext = 'png'
                    elif 'jpeg' in header or 'jpg' in header:
                        ext = 'jpg'
                    elif 'gif' in header:
                        ext = 'gif'
                    elif 'webp' in header:
                        ext = 'webp'
                    else:
                        ext = 'png'  # 默认PNG
                else:
                    data = image_data
                    ext = 'png'
                
                # 创建文件名
                filename = f"{chrome_id}_thumb.{ext}"
                filepath = os.path.join('files', filename)
                
                # 保存文件
                with open(filepath, 'wb') as f:
                    f.write(base64.b64decode(data))
                
                # 更新数据库记录，将image_data字段清空
                cursor.execute("UPDATE prompts SET image_data = NULL WHERE id = ?", (record_id,))
                
                migrated_count += 1
                print(f"✓ 迁移记录 {record_id} (Chrome ID: {chrome_id}) -> {filename}")
                
            except Exception as e:
                error_count += 1
                print(f"✗ 迁移记录 {record_id} 失败: {e}")
        
        # 提交更改
        conn.commit()
        
        print(f"\n迁移完成:")
        print(f"  成功迁移: {migrated_count} 条记录")
        print(f"  失败: {error_count} 条记录")
        
        # 检查数据库大小变化
        db_size = os.path.getsize(db_path)
        print(f"  数据库文件大小: {db_size / 1024 / 1024:.2f} MB")
        
    except Exception as e:
        print(f"迁移过程中发生错误: {e}")
        conn.rollback()
    finally:
        conn.close()

def check_migration_result():
    """检查迁移结果"""
    db_path = 'instance/prompt_manager.db'
    
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    
    try:
        # 检查还有多少记录包含图片数据
        cursor.execute("SELECT COUNT(*) FROM prompts WHERE image_data IS NOT NULL AND image_data != ''")
        remaining_count = cursor.fetchone()[0]
        
        # 检查files目录中的文件数量
        files_count = len([f for f in os.listdir('files') if f.endswith(('.png', '.jpg', '.jpeg', '.gif', '.webp'))])
        
        print(f"\n迁移结果检查:")
        print(f"  数据库中仍有图片数据的记录: {remaining_count}")
        print(f"  files目录中的图片文件: {files_count}")
        
        if remaining_count == 0:
            print("✓ 所有图片数据已成功迁移")
        else:
            print("⚠ 仍有图片数据未迁移")
            
    finally:
        conn.close()

if __name__ == '__main__':
    print("=" * 50)
    print("图片数据迁移脚本")
    print("=" * 50)
    
    # 备份数据库
    db_path = 'instance/prompt_manager.db'
    backup_path = f'instance/prompt_manager_backup_{datetime.now().strftime("%Y%m%d_%H%M%S")}.db'
    
    if os.path.exists(db_path):
        import shutil
        shutil.copy2(db_path, backup_path)
        print(f"已创建数据库备份: {backup_path}")
    
    # 执行迁移
    migrate_images()
    
    # 检查结果
    check_migration_result()
    
    print("\n迁移完成！")

